package io.github.handyang.algorithms.search;

import lombok.Getter;
import lombok.Setter;
import lombok.ToString;
import lombok.extern.slf4j.Slf4j;
import org.junit.Test;

import java.util.LinkedList;
import java.util.List;

@Slf4j
public class RedBlackBST<K extends Comparable<K>, V> implements SymbolTable<K, V> {

    @Getter
    private Node<K, V> root;
    private int size;

    @Override
    public void put(K key, V value) {
        log.debug("key:{},value:{}", key, value);

        root = put0(root, key, value);
        root.color = Node.BLACK;
    }

    private Node<K, V> put0(Node<K, V> root, K key, V value) {
        if (root == null) {
            size++;
            return new Node<>(key, value, Node.RED);

        }

        int compare = root.getKey().compareTo(key);
        if (compare == 0) root.setValue(value);
        else if (compare > 0) root.setLeft(put0((Node<K, V>) root.getLeft(), key, value));
        else root.setRight(put0((Node<K, V>) root.getRight(), key, value));

        if (!isRed(root.getLeft()) && isRed(root.getRight())) root = rotateLeft(root);
        if (isRed(root.getLeft()) && isRed(root.getLeft().getLeft())) root = rotateRight(root);
        if (isRed(root.getLeft()) && isRed(root.getRight())) flipColors(root);

        return root;
    }

    private Node<K, V> rotateLeft(Node<K, V> node) {
        log.debug("rotateLeft:{}", node);
        Node<K, V> right = (Node<K, V>) node.getRight();

        node.setRight(right.getLeft());

        right.setLeft(node);
        right.setColor(node.color);
        node.setColor(Node.RED);

        return right;
    }

    private Node<K, V> rotateRight(Node<K, V> node) {
        log.debug("rotateRight:{}", node);

        Node<K, V> left = (Node<K, V>) node.getLeft();

        node.setLeft(left.getRight());

        left.setRight(node);
        left.setColor(node.color);
        node.setColor(Node.RED);

        return left;
    }

    private void flipColors(Node<K, V> node) {
        log.debug("flipColors:{}", node);

        node.setColor(Node.RED);
        ((Node) node.getLeft()).setColor(Node.BLACK);
        ((Node) node.getRight()).setColor(Node.BLACK);
    }

    private boolean isRed(TreeNode<K, V> node) {
        if (node == null) return false;

        return ((Node) node).color == Node.RED;
    }

    @Override
    public V get(K key) {
        return get0(root, key);
    }

    private V get0(TreeNode<K, V> root, K key) {
        if (root == null) return null;

        int compare = root.getKey().compareTo(key);
        if (compare > 0) return get0(root.getLeft(), key);
        else if (compare < 0) return get0(root.getRight(), key);
        else return root.getValue();
    }

    @Override
    public int size() {
        return size;
    }

    @Override
    public Iterable<K> keys() {
        List<K> keys = new LinkedList<>();
        allKey(keys, this.root);

        return keys;
    }

    private void allKey(List<K> keys, TreeNode<K, V> root) {
        if (root == null) return;

        keys.add(root.getKey());

        allKey(keys, root.getLeft());
        allKey(keys, root.getRight());
    }

    @Override
    public void clear() {
        this.root = null;
        this.size = 0;
    }

    @Setter
    @ToString(callSuper = true)
    private static final class Node<K extends Comparable<K>, V> extends TreeNode<K, V> {
        private static final int RED = 1;
        private static final int BLACK = 2;

        //父节点指向该节点的颜色
        private int color;

        public Node(K key, V value, int color) {
            super(key, value);
            this.color = color;
        }
    }


    @Test
    public void testRotateLeft() {
        Node<String, Integer> root = new Node<>("15", 1, Node.BLACK);
        root.setLeft(new Node<>("11", 1, Node.BLACK));
        root.setRight(new Node<>("17", 1, Node.RED));

        root.getRight().setLeft(new Node<>("16", 1, Node.RED));
        root.getRight().setRight(new Node<>("20", 1, Node.BLACK));

        root = (Node<String, Integer>) rotateLeft((Node<K, V>) root);

        System.out.println(root);
    }

    @Test
    public void testRotateRight() {
        Node<String, Integer> root = new Node<>("15", 1, Node.BLACK);
        root.setLeft(new Node<>("11", 1, Node.RED));
        root.setRight(new Node<>("17", 1, Node.BLACK));

        root.getLeft().setLeft(new Node<>("6", 1, Node.RED));
        root.getLeft().setRight(new Node<>("13", 1, Node.BLACK));

        root = (Node<String, Integer>) rotateRight((Node<K, V>) root);

        System.out.println(root);
    }
}
